import os
import json
import random
import base64
from io import BytesIO
from typing import Optional, List, Tuple
import uvicorn
from fastapi import FastAPI, File, Form, UploadFile
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import torch
from PIL import Image
from transformers import (
    AutoProcessor,
    Qwen2_5_VLForConditionalGeneration,
)

app = FastAPI(title="商品类别预测 API")

# Configure CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Load model and processor
MODEL_PATH = "./output/tari-product-image"
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    MODEL_PATH, device_map="auto"
)
processor = AutoProcessor.from_pretrained(MODEL_PATH)

# Load product data
with open("images/data.json", "r", encoding="utf-8") as f:
    product_data = json.load(f)

# Create static files directory if it doesn't exist
os.makedirs("static", exist_ok=True)
os.makedirs("static/img", exist_ok=True)

# Mount static files directory
app.mount("/static", StaticFiles(directory="static"), name="static")

# Mount images directory
app.mount("/images", StaticFiles(directory="images"), name="images")

class ProductResponse(BaseModel):
    category: str
    similar_products: List[dict]

@app.get("/", response_class=HTMLResponse)
async def get_html():
    with open("static/index.html", "r", encoding="utf-8") as f:
        html_content = f.read()
    return html_content

@app.post("/predict", response_model=ProductResponse)
async def predict(
    file: Optional[UploadFile] = File(None),
    text: Optional[str] = Form(None)
):
    # Check if both inputs are missing
    if file is None and (text is None or text.strip() == ""):
        return {"error": "请提供图片或文本输入"}
    
    # Process image if provided
    if file is not None:
        contents = await file.read()
        image = Image.open(BytesIO(contents)).convert("RGB")
        
        # Image classification
        processed_conversation = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "image": image,
                    },
                    {
                        "type": "text",
                        "text": "这张图片展示的是什么商品类别？直接输出类别名称，不要输出任何解释。",
                    },
                ],
            }
        ]
    # Process text if no image but text is provided
    elif text and text.strip():
        processed_conversation = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": f"商品标题：\"{text}\"\n请根据这个标题判断商品的类别。直接输出类别名称，不要输出任何解释。",
                    },
                ],
            }
        ]
    
    # Get prediction
    predicted_category = predict_category(processed_conversation)
    
    # Get all products from this category for analytics
    category_products = [
        item for item in product_data.values() if item["label"] == predicted_category
    ]
    
    # Generate analytics data from all category products
    analytics = generate_analytics(category_products)
    
    # Instead of limiting to 5 random products, send all products
    # This allows the frontend to handle pagination
    similar_products = prepare_similar_products(category_products)
    
    return {
        "category": predicted_category,
        "similar_products": similar_products,
        "analytics": analytics
    }

def predict_category(processed_conversation):
    inputs = processor.apply_chat_template(
        processed_conversation,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device)

    output_ids = model.generate(**inputs, max_new_tokens=128)
    generated_ids = [
        output_ids[i][len(inputs.input_ids[i]) :] for i in range(len(output_ids))
    ]
    predicted_category = processor.batch_decode(
        generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
    )[0].strip()
    
    return predicted_category

def prepare_similar_products(products):
    similar_products = []
    for product in products:
        img_path = f"/images/img/{os.path.basename(product['url'])}"
        
        # 获取或生成附加字段
        price = product.get("price", random.randint(1, 500))
        platform = product.get("platform", random.choice(["京东", "拼多多", "抖音商城", "天猫国际", "药房网商城", "淘宝", "携程", "唯品会", "考拉海购", "妈妈网孕育"]))
        comment = product.get("comment", random.choice(["积极", "消极", "中立"]))
        
        # 创建包含所有原始字段的产品数据
        product_data = dict(product)  # 复制原始产品的所有字段
        
        # 更新或添加特定字段
        product_data.update({
            "image_url": img_path,
            "price": price,
            "platform": platform,
            "comment": comment
        })
        
        similar_products.append(product_data)
    
    return similar_products

def generate_analytics(category_products):
    # Platform distribution
    platforms = {}
    for product in category_products:
        platform = product.get("platform", "未知")
        if platform in platforms:
            platforms[platform] += 1
        else:
            platforms[platform] = 1
    
    # 动态生成价格区间
    price_ranges = dynamic_price_ranges(category_products)
    
    # Comment ratio
    comments = {
        "积极": 0,
        "中立": 0,
        "消极": 0
    }
    
    for product in category_products:
        comment = product.get("comment", random.choice(["积极", "消极", "中立"]))
        if comment in comments:
            comments[comment] += 1
    
    return {
        "platforms": platforms,
        "priceRanges": price_ranges,
        "comments": comments
    }

def dynamic_price_ranges(products):
    # 提取所有价格
    prices = []
    for product in products:
        price = product.get("price")
        if isinstance(price, str):
            try:
                price = float(price)
                prices.append(price)
            except:
                pass
        elif isinstance(price, (int, float)):
            prices.append(price)
    
    # 如果没有足够的价格数据，使用默认区间
    if len(prices) < 10:
        return default_price_ranges()
    
    # 分析价格分布
    min_price = min(prices)
    max_price = max(prices)
    
    # 如果价格范围太小，使用固定间隔
    if max_price - min_price < 100:
        return simple_price_ranges(min_price, max_price)
    
    # 使用更复杂的分布分析方法
    return optimized_price_ranges(prices)

def default_price_ranges():
    # 默认价格区间
    return {
        "[0-50]": 0,
        "[50-100]": 0,
        "[100-150]": 0,
        "[150-200]": 0,
        "[200-250]": 0,
        "[250-300]": 0,
        "[300-350]": 0,
        "[350-400]": 0,
        "[400-450]": 0,
        "[450-500]": 0
    }

def simple_price_ranges(min_price, max_price):
    # 简单的等间隔价格区间
    step = (max_price - min_price) / 10
    ranges = {}
    
    for i in range(10):
        start = int(min_price + i * step)
        end = int(min_price + (i + 1) * step)
        if i == 9:  # 确保最后一个区间包含最大值
            end = int(max_price) + 1
        ranges[f"[{start}-{end}]"] = 0
    
    return ranges

def optimized_price_ranges(prices):
    # 将价格按分位数分布，以获得更均匀的分布
    prices.sort()
    n = len(prices)
    ranges = {}
    
    # 分成10个区间，每个区间包含大约相同数量的数据点
    for i in range(10):
        start_idx = (i * n) // 10
        end_idx = ((i + 1) * n) // 10 - 1
        if end_idx < 0 or start_idx >= n:
            continue
            
        start = int(prices[start_idx])
        end = int(prices[min(end_idx, n-1)])
        
        # 确保区间不重叠
        if i > 0 and start <= prev_end:
            start = prev_end + 1
        
        # 确保最后一个区间包含最大值
        if i == 9:
            end = int(max(prices)) + 1
            
        # 确保区间有效
        if start < end:
            ranges[f"[{start}-{end}]"] = 0
            prev_end = end
        
    # 填充每个区间的商品数量
    for price in prices:
        for range_str in ranges:
            match = range_str.replace('[', '').replace(']', '').split('-')
            range_start = int(match[0])
            range_end = int(match[1])
            
            if range_start <= price < range_end or (price == range_end and range_str.endswith(f"{range_end}]")):
                ranges[range_str] += 1
                break
    
    return ranges

if __name__ == "__main__":
    uvicorn.run("api:app", host="0.0.0.0", port=8000, reload=True) 